package com.leonzhangxf.client.register;

import com.leonzhangxf.client.loadbalance.LoadBalancer;
import com.leonzhangxf.client.loadbalance.LoadDataCleaner;
import lombok.Value;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.api.CuratorWatcher;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.data.Stat;
import org.springframework.util.StringUtils;

import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author leonzhangxf
 */
@Slf4j
public class ZkServiceLoader implements ServiceLoader, LoadBalancer, LoadDataCleaner {

    private final CuratorFramework client;

    ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(3, new ThreadFactory() {

        private final AtomicInteger count = new AtomicInteger();

        @Override
        public Thread newThread(Runnable r) {
            return new Thread(r, "LoadDataCleaner-" + count.getAndIncrement());
        }
    });

    public ZkServiceLoader(String connectString) {
        client = CuratorFrameworkFactory.builder()
            .connectString(connectString)
            .retryPolicy(new ExponentialBackoffRetry(1000, 3, 10000))
            .namespace("program_2").build();

        client.start();
    }

    @Override
    public List<Node> loadNode(String serviceId, NodeListChangeCallback callback) throws Exception {
        String serviceIdNode = getServiceIdNode(serviceId);
        Stat stat = client.checkExists()
            .creatingParentsIfNeeded()
            .forPath(serviceIdNode);
        if (null == stat) {
            return Collections.emptyList();
        }

        client.getChildren().usingWatcher(new ChildCuratorWatcher(client, serviceIdNode, callback)).forPath(serviceIdNode);

        List<String> children = client.getChildren().forPath(serviceIdNode);
        return toNodes(children);
    }

    private static String getServiceIdNode(String serviceId) {
        return String.format("/%s", serviceId);
    }

    private static String nodeToPath(Node node) {
        return String.format("%s_%d", node.getIp(), node.getPort());
    }

    private static String getNodePath(String serviceId, String nodePath) {
        return String.format("/%s/%s", serviceId, nodePath);
    }

    static List<Node> toNodes(List<String> children) {
        return Optional.ofNullable(children).orElse(Collections.emptyList())
            .stream().map(nodePathToNodeMapping()).collect(Collectors.toList());
    }

    private static Function<String, Node> nodePathToNodeMapping() {
        return raw -> {
            String[] str = raw.split("_");
            return new Node(str[0], Integer.parseInt(str[1]));
        };
    }

    private static LoadData getLoadData(String data) {
        if (!StringUtils.hasText(data)) {
            return null;
        }
        try {
            String[] split = data.split(",");
            long costTimeMillis = Long.parseLong(split[0]);
            long timestamp = Long.parseLong(split[1]);
            return new LoadData(costTimeMillis, timestamp);
        } catch (Exception ex) {
            return null;
        }
    }

    /**
     * 记录服务请求响应时间
     *
     * @param serviceId      服务Id
     * @param node           节点
     * @param costTimeMillis 耗时
     */
    public void updateLoadData(String serviceId, Node node, long costTimeMillis) {
        String nodePath = getNodePath(serviceId, nodeToPath(node));
        try {
            if (null != client.checkExists().forPath(nodePath)) {
                log.info("update node: {} load data: {}", node, costTimeMillis);
                client.setData().forPath(
                    nodePath,
                    String.format("%d,%d", costTimeMillis, System.currentTimeMillis()).getBytes()
                );
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public Node choose(String serviceId) {
        // 找到耗时最少的一个节点
        List<String> children;
        try {
            children = client.getChildren().forPath(getServiceIdNode(serviceId));
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
        Node resultNode = children.stream().map(node -> {
            byte[] bytes;
            String nodePath = getNodePath(serviceId, node);
            try {
                bytes = client.getData().forPath(nodePath);
            } catch (Exception e) {
                e.printStackTrace();
                return null;
            }
            String data = new String(bytes);
            log.info("raw node: {} load data: {}", nodePath, data);
            LoadData loadData = getLoadData(data);
            if (null == loadData) {
                loadData = new LoadData(0L, System.currentTimeMillis());
            }
            Node nodeVal = Optional.of(node).map(nodePathToNodeMapping()).orElse(null);
            if (null == nodeVal) {
                return null;
            }
            log.info("node: {}, loadData: {}", nodeVal, loadData);
            return new NodeWrapper(nodeVal, loadData);
        }).min((n1, n2) -> {
            long n1Cost = Optional.ofNullable(n1).map(NodeWrapper::getLoadData)
                .map(LoadData::getCostTimeMillis).orElse(0L);
            long n2Cost = Optional.ofNullable(n2).map(NodeWrapper::getLoadData)
                .map(LoadData::getCostTimeMillis).orElse(0L);
            return Long.compare(n1Cost, n2Cost);
        }).map(NodeWrapper::getNode).orElse(null);
        log.info("choose node: {}", resultNode);
        return resultNode;
    }

    @Value
    static class NodeWrapper {
        Node node;
        LoadData loadData;
    }

    @Value
    static class LoadData {
        Long costTimeMillis;
        Long timestamp;
    }

    @Override
    public void scheduleCleanData(String serviceId) {
        executor.scheduleAtFixedRate(new DataCleaner(serviceId, client),
            2000, 5000, TimeUnit.MILLISECONDS);
    }

    @Slf4j
    public static class DataCleaner implements Runnable {

        private final String serviceId;

        private final CuratorFramework client;

        public DataCleaner(String serviceId, CuratorFramework client) {
            this.serviceId = serviceId;
            this.client = client;
        }

        @Override
        public void run() {
            log.info("清理serviceId:{} 的超时负载数据", serviceId);
            // 清理对应服务的节点负载数据
            List<String> children = null;
            try {
                children = client.getChildren().forPath(getServiceIdNode(serviceId));
            } catch (Exception e) {
                e.printStackTrace();
            }
            if (null == children || children.isEmpty()) {
                return;
            }
            for (String node : children) {
                String nodePath = getNodePath(serviceId, node);
                byte[] bytes = new byte[0];
                try {
                    bytes = client.getData().forPath(nodePath);
                } catch (Exception e) {
                    e.printStackTrace();
                }
                LoadData loadData = getLoadData(new String(bytes));
                if (null == loadData) {
                    continue;
                }
                long currentTimeMillis = System.currentTimeMillis();
                if (currentTimeMillis - loadData.getTimestamp() > 5000) {
                    // 时间大于5秒清理
                    try {
                        if (null != client.checkExists().forPath(nodePath)) {
                            client.setData().forPath(nodePath, "".getBytes());
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
            }
        }
    }

    public static class ChildCuratorWatcher implements CuratorWatcher {

        private final CuratorFramework client;

        private final String serviceIdNode;

        private final NodeListChangeCallback callback;

        public ChildCuratorWatcher(CuratorFramework client, String serviceIdNode, NodeListChangeCallback callback) {
            this.client = client;
            this.serviceIdNode = serviceIdNode;
            this.callback = callback;
        }

        @Override
        public void process(WatchedEvent watchedEvent) throws Exception {
            if (watchedEvent.getState() == Watcher.Event.KeeperState.SyncConnected
                && watchedEvent.getType() == Watcher.Event.EventType.NodeChildrenChanged) {
                List<String> children = client.getChildren().forPath(serviceIdNode);
                System.out.println("新的子节点信息" + children);

                if (null != callback) {
                    callback.nodeListChanged(toNodes(children));
                }

                client.getChildren().usingWatcher(this).forPath(serviceIdNode);
            }
        }
    }
}
